package com.example;

import javafx.application.Application;
import javafx.application.Platform;
import javafx.collections.FXCollections;
import javafx.collections.ObservableList;
import javafx.geometry.Insets;
import javafx.geometry.Pos;
import javafx.scene.Scene;
import javafx.scene.control.*;
import javafx.scene.layout.*;
import javafx.stage.FileChooser;
import javafx.stage.Stage;
import org.pcap4j.core.*;
import org.pcap4j.packet.*;

import java.io.File;
import java.io.FileNotFoundException;
import java.io.PrintWriter;
import java.net.Inet4Address;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Queue;
import java.util.Set;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.stream.Collectors;

public class PacketSniffer extends Application {

    private PacketSnifferUI ui;
    private PacketSnifferController controller;

    @Override
    public void start(Stage primaryStage) {
        ui = new PacketSnifferUI(primaryStage);
        controller = new PacketSnifferController(ui);
        controller.initialize();
    }

    public static void main(String[] args) {
        launch(args);
    }
}

class PacketSnifferUI {

    private Stage primaryStage;
    private ComboBox<String> interfaceCombo;
    private ComboBox<String> protocolCombo;
    private TextField ipField;
    private Button startButton;
    private Button stopButton;
    private TextArea textArea;

    public PacketSnifferUI(Stage primaryStage) {
        this.primaryStage = primaryStage;
        setupUI();
    }

    private void setupUI() {
        primaryStage.setTitle("Packet Sniffer");
        primaryStage.setOnCloseRequest(event -> {
            Platform.exit();
            System.exit(0);
        });

        // North Panel: Interface selection, IP, and Protocol input
        GridPane northPanel = new GridPane();
        northPanel.setPadding(new Insets(20));
        northPanel.setHgap(20);
        northPanel.setVgap(20);
        northPanel.setAlignment(Pos.CENTER);

        northPanel.add(new Label("Select Interface:"), 0, 0);
        interfaceCombo = new ComboBox<>();
        northPanel.add(interfaceCombo, 1, 0);

        northPanel.add(new Label("Enter IP:"), 0, 1);
        ipField = new TextField();
        ipField.setPromptText("Enter target IP address here");
        northPanel.add(ipField, 1, 1);

        northPanel.add(new Label("Select Protocol:"), 0, 2);
        protocolCombo = new ComboBox<>(FXCollections.observableArrayList("ALL", "TCP", "UDP", "HTTP", "HTTPS"));
        northPanel.add(protocolCombo, 1, 2);

        // Center Panel: Text area for displaying packet info
        textArea = new TextArea();
        textArea.setEditable(false);
        textArea.setWrapText(true);
        textArea.setMinHeight(500);
        VBox.setVgrow(textArea, Priority.ALWAYS);

        // South Panel: Start and Stop buttons
        HBox southPanel = new HBox(10);
        southPanel.setPadding(new Insets(20));
        southPanel.setAlignment(Pos.CENTER);
        startButton = new Button("Start Capturing");
        stopButton = new Button("Stop Capturing");
        Button clearButton = new Button("Clear Output");
        Button exportButton = new Button("Export Output");
        stopButton.setDisable(true);
        southPanel.getChildren().addAll(startButton, stopButton, clearButton, exportButton);

        // Layout all panels
        VBox layout = new VBox(10);
        layout.getChildren().addAll(northPanel, textArea, southPanel);
        layout.setPadding(new Insets(10));

        // Set Scene and Stage
        Scene scene = new Scene(layout, 800, 600);
        scene.getStylesheets().add(getClass().getResource("/style.css").toExternalForm());
        primaryStage.setScene(scene);
        primaryStage.setMaximized(true);
        primaryStage.show();

        // Add button actions
        clearButton.setOnAction(e -> textArea.clear());
        exportButton.setOnAction(e -> exportOutput());
    }

    public ComboBox<String> getInterfaceCombo() {
        return interfaceCombo;
    }

    public ComboBox<String> getProtocolCombo() {
        return protocolCombo;
    }

    public TextField getIpField() {
        return ipField;
    }

    public Button getStartButton() {
        return startButton;
    }

    public Button getStopButton() {
        return stopButton;
    }

    public TextArea getTextArea() {
        return textArea;
    }

    private void exportOutput() {
        FileChooser fileChooser = new FileChooser();
        fileChooser.setTitle("Save Output");
        fileChooser.getExtensionFilters().addAll(new FileChooser.ExtensionFilter("Text Files", "*.txt"));
        File file = fileChooser.showSaveDialog(primaryStage);
        if (file != null) {
            try (PrintWriter writer = new PrintWriter(file)) {
                writer.write(textArea.getText());
            } catch (FileNotFoundException e) {
                new Alert(Alert.AlertType.ERROR, "Failed to save output: " + e.getMessage(), ButtonType.OK).showAndWait();
            }
        }
    }
}

class PacketSnifferController {

    private PacketSnifferUI ui;
    private PacketSnifferService service;

    public PacketSnifferController(PacketSnifferUI ui) {
        this.ui = ui;
        this.service = new PacketSnifferService(ui);
    }

    public void initialize() {
        setupInterfaceList();
        ui.getStartButton().setOnAction(e -> service.startCapturing(0));
        ui.getStopButton().setOnAction(e -> service.stopCapturing());
        ui.getTextArea().setWrapText(true);
    }

    private void setupInterfaceList() {
        try {
            List<PcapNetworkInterface> allDevs = Pcaps.findAllDevs();
            ObservableList<String> interfaces = FXCollections.observableArrayList();
            for (PcapNetworkInterface dev : allDevs) {
                List<PcapAddress> addresses = dev.getAddresses();
                String ipInfo = addresses.stream()
                        .filter(addr -> addr.getAddress() instanceof Inet4Address)
                        .map(addr -> "IP: " + addr.getAddress().getHostAddress() + ", Mask: " + addr.getNetmask().getHostAddress())
                        .collect(Collectors.joining(", "));

                if (ipInfo.isEmpty()) {
                    ipInfo = "No IPv4 address";
                }
                interfaces.add(String.format("%s (%s) [%s]", dev.getName(), dev.getDescription(), ipInfo));
            }
            ui.getInterfaceCombo().setItems(interfaces);
        } catch (PcapNativeException e) {
            alertError("Failed to list network interfaces: " + e.getMessage());
        }
    }

    private void alertError(String message) {
        Platform.runLater(() -> {
            Alert alert = new Alert(Alert.AlertType.ERROR, message, ButtonType.OK);
            alert.showAndWait();
        });
    }
}

class PacketSnifferService {

    private PacketSnifferUI ui;
    private ExecutorService pool;
    private Set<PcapHandle> activeHandles = ConcurrentHashMap.newKeySet();
    private AtomicBoolean errorDisplayed = new AtomicBoolean(false);
    private Queue<String> packetInfoQueue = new ConcurrentLinkedQueue<>();

    public PacketSnifferService(PacketSnifferUI ui) {
        this.ui = ui;
        setupPacketInfoDisplay();
    }

    public void startCapturing(int a) {
        try {
            Thread.sleep(100);
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
        if(a==0)startCapturing(1);
        int index = ui.getInterfaceCombo().getSelectionModel().getSelectedIndex();
        String ip = ui.getIpField().getText();
        String selectedProtocol = ui.getProtocolCombo().getValue();

        if (index == -1 || ip.isEmpty()) {
            alertError("Please select an interface and enter a valid IP address.");
            return;
        }

        try {
            PcapNetworkInterface nif = Pcaps.findAllDevs().get(index);
            String filter = createFilter(selectedProtocol, ip);
            int numberOfThreads = 4;
            pool = Executors.newFixedThreadPool(numberOfThreads);

            for (int i = 0; i < numberOfThreads; i++) {
                final PcapHandle threadHandle = nif.openLive(65536, PcapNetworkInterface.PromiscuousMode.PROMISCUOUS, 5);
                threadHandle.setFilter(filter, BpfProgram.BpfCompileMode.OPTIMIZE);
                activeHandles.add(threadHandle);
                pool.submit(() -> {
                    try {
                        capturePackets(threadHandle);
                    } finally {
                        threadHandle.close();
                        activeHandles.remove(threadHandle);
                    }
                });
            }

            ui.getStartButton().setDisable(true);
            ui.getStopButton().setDisable(false);
        } catch (Exception e) {
            alertError("Error setting up capture: " + e.getMessage());
        }
    }

    private void capturePackets(PcapHandle handle) {
        try {
            PacketListener listener = packet -> {
                if (Thread.currentThread().isInterrupted()) {
                    try {
                        handle.breakLoop();
                    } catch (NotOpenException e) {
                        throw new RuntimeException(e);
                    }
                    return;
                }
                displayPacket(packet, ui.getProtocolCombo().getValue());
            };
            handle.loop(-1, listener);
        } catch (PcapNativeException | InterruptedException | NotOpenException e) {
            if (errorDisplayed.compareAndSet(false, true)) {
                alertError("Error capturing packets: " + e.getMessage());
            }
        } finally {
            handle.close();
        }
    }

    private void displayPacket(Packet packet, String protocol) {
        StringBuilder sb = new StringBuilder();
        if (packet.contains(EthernetPacket.class)) {
            EthernetPacket ethernetPacket = packet.get(EthernetPacket.class);
            EthernetPacket.EthernetHeader ethernetHeader = ethernetPacket.getHeader();
            sb.append(String.format("Ethernet Header -> Dst MAC: %s, Src MAC: %s, Type: %s\n",
                    ethernetHeader.getDstAddr(), ethernetHeader.getSrcAddr(), ethernetHeader.getType()));
        }
        if (packet.contains(IpV4Packet.class)) {
            IpV4Packet ipV4Packet = packet.get(IpV4Packet.class);
            IpV4Packet.IpV4Header ipHeader = ipV4Packet.getHeader();
            sb.append(String.format("IP Header -> Version: %s, Header Length: %d, Total Length: %d, Identification: %d, TTL: %d, Protocol: %s, Src Addr: %s, Dst Addr: %s\n",
                    ipHeader.getVersion(), ipHeader.getIhl() * 4, ipHeader.getTotalLengthAsInt(),
                    ipHeader.getIdentificationAsInt(), ipHeader.getTtlAsInt(), ipHeader.getProtocol(),
                    ipHeader.getSrcAddr(), ipHeader.getDstAddr()));
        }
        if (packet.contains(TcpPacket.class)) {
            TcpPacket tcpPacket = packet.get(TcpPacket.class);
            TcpPacket.TcpHeader tcpHeader = tcpPacket.getHeader();
            sb.append(String.format("TCP Header -> Src Port: %s, Dst Port: %s, Seq Number: %d, Ack Number: %d, Data Offset: %d, Window: %d, Checksum: %d, Urgent Pointer: %d\n",
                    tcpHeader.getSrcPort(), tcpHeader.getDstPort(), tcpHeader.getSequenceNumber(),
                    tcpHeader.getAcknowledgmentNumber(), tcpHeader.getDataOffsetAsInt() * 4,
                    tcpHeader.getWindowAsInt(), tcpHeader.getChecksum(), tcpHeader.getUrgentPointerAsInt()));

            // 获取并显示 TCP 数据载荷
            if (tcpPacket.getPayload() != null) {
                byte[] tcpData = tcpPacket.getPayload().getRawData();
                if (tcpData.length > 0) {
                    try {
                        String payloadData = new String(tcpData, StandardCharsets.UTF_8); // 尝试使用 UTF-8 编码解码数据
                        // 检测解码后的字符串是否包含乱码
                        if (!looksLikeBinary(payloadData)) {
                            sb.append("TCP Payload (text): ").append(payloadData).append('\n');
                        } else {
                            // 如果出现乱码，显示十六进制形式
                            String hexData = bytesToHex(tcpData);
                            sb.append("TCP Payload (hex): ").append(hexData).append('\n');
                        }
                    } catch (Exception e) {
                        sb.append("TCP Payload: [Error decoding payload]").append('\n');
                    }
                } else {
                    sb.append("TCP Payload: <none>").append('\n');  // 显示无数据载荷信息
                }
            } else {
                sb.append("TCP Payload: <none>").append('\n');  // 显示无数据载荷信息
            }
        }
        if (packet.contains(UdpPacket.class)) {
            UdpPacket udpPacket = packet.get(UdpPacket.class);
            UdpPacket.UdpHeader udpHeader = udpPacket.getHeader();
            sb.append(String.format("UDP Header -> Src Port: %s, Dst Port: %s, Length: %d, Checksum: %d\n",
                    udpHeader.getSrcPort(), udpHeader.getDstPort(), udpHeader.getLengthAsInt(),
                    udpHeader.getChecksum()));
            // 获取并显示 UDP 数据载荷
            if (udpPacket.getPayload() != null) {
                byte[] udpData = udpPacket.getPayload().getRawData();
                if (udpData.length > 0) {
                    try {
                        String payloadData = new String(udpData, StandardCharsets.UTF_8); // 尝试使用 UTF-8 编码解码数据
                        // 检测解码后的字符串是否可能包含乱码
                        if (!looksLikeBinary(payloadData)) {
                            sb.append("UDP Payload (text): ").append(payloadData).append('\n');
                        } else {
                            // 如果检测到乱码，显示十六进制形式
                            String hexData = bytesToHex(udpData);
                            sb.append("UDP Payload (hex): ").append(hexData).append('\n');
                        }
                    } catch (Exception e) {
                        sb.append("UDP Payload: [Error decoding payload]").append('\n');
                    }
                } else {
                    sb.append("UDP Payload: <none>").append('\n');  // 显示无数据载荷信息
                }
            } else {
                sb.append("UDP Payload: <none>").append('\n');  // 显示无数据载荷信息
            }
        }
        sb.append('\n');
        Platform.runLater(() -> ui.getTextArea().appendText(sb.toString()));
    }

    private void setupPacketInfoDisplay() {
        ScheduledExecutorService displayExecutor = Executors.newSingleThreadScheduledExecutor();
        displayExecutor.scheduleAtFixedRate(() -> {
            StringBuilder sb = new StringBuilder();
            String packetInfo;
            while ((packetInfo = packetInfoQueue.poll()) != null) {
                sb.append(packetInfo);
            }
            if (sb.length() > 0) {
                Platform.runLater(() -> ui.getTextArea().appendText(sb.toString()));
            }
        }, 0, 1000, TimeUnit.SECONDS);
    }

    public void stopCapturing() {
        if (pool != null) {
            pool.shutdownNow();
            for (PcapHandle handle : activeHandles) {
                try {
                    handle.breakLoop();
                } catch (NotOpenException e) {
                    alertError("Error stopping capture: " + e.getMessage());
                }
            }
        }
        activeHandles.clear();
        ui.getStartButton().setDisable(false);
        ui.getStopButton().setDisable(true);
    }

    private String createFilter(String protocol, String ip) {
        String protocolFilter = "";
        switch (protocol) {
            case "TCP":
                protocolFilter = "tcp";
                break;
            case "UDP":
                protocolFilter = "udp";
                break;
            case "HTTP":
                protocolFilter = "tcp port 80";
                break;
            case "HTTPS":
                protocolFilter = "tcp port 443";
                break;
            case "ALL":
                protocolFilter = "";
                break;
        }
        return String.format("%s%s", protocolFilter.isEmpty() ? "" : protocolFilter + " and ", "(src host " + ip + " or dst host " + ip + ")");
    }


    private void alertError(String message) {
        Platform.runLater(() -> {
            Alert alert = new Alert(Alert.AlertType.ERROR, message, ButtonType.OK);
            alert.showAndWait();
        });
    }

    private boolean looksLikeBinary(String data) {
        for (char c : data.toCharArray()) {
            if (Character.isISOControl(c) && !Character.isWhitespace(c)) {
                return true;
            }
        }
        return false;
    }

    private String bytesToHex(byte[] bytes) {
        StringBuilder sb = new StringBuilder();
        for (byte b : bytes) {
            sb.append(String.format("%02X ", b));
        }
        return sb.toString();
    }
}
